import cv2
import numpy as np
import matplotlib.pyplot as plt

def embed_watermark(carrier_path, watermark_path, alpha=0.3, output_path="embedded_image.png"):
    # 读取载体图像和水印图像
    carrier = cv2.imread(carrier_path, cv2.IMREAD_GRAYSCALE)
    watermark = cv2.imread(watermark_path, cv2.IMREAD_GRAYSCALE)

    # 调整水印图像大小与载体图像一致
    watermark_resized = cv2.resize(watermark, (carrier.shape[1], carrier.shape[0]))

    # 计算DFT
    carrier_dft = np.fft.fftshift(np.fft.fft2(carrier))
    watermark_dft = np.fft.fftshift(np.fft.fft2(watermark_resized))

    # 嵌入水印 (频域操作)
    embedded_dft = carrier_dft + alpha * watermark_dft

    # 逆DFT还原图像
    embedded_dft = np.fft.ifftshift(embedded_dft)
    embedded_dft = np.fft.ifft2(embedded_dft)
    reconstructed_image = np.abs(embedded_dft)

    # 保存结果
    plt.imsave(output_path, reconstructed_image, cmap='gray')
    print(f"嵌入水印的图像已保存到 {output_path}")

# 调用
embed_watermark("alena.bmp", "dog.png")

